package llutyo.andcrypt;

/*Copyright 2012 Akos Ludanyi

		   Licensed under the Apache License, Version 2.0 (the "License");
		   you may not use this file except in compliance with the License.
		   You may obtain a copy of the License at

		       http://www.apache.org/licenses/LICENSE-2.0

		   Unless required by applicable law or agreed to in writing, software
		   distributed under the License is distributed on an "AS IS" BASIS,
		   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
		   See the License for the specific language governing permissions and
		   limitations under the License.
*/
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.RandomAccessFile;
import java.security.MessageDigest;
import java.security.spec.AlgorithmParameterSpec;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import android.app.Activity;
import android.app.AlertDialog;
import android.app.Dialog;
import android.app.ProgressDialog;
import android.content.DialogInterface;
import android.content.Intent;
import android.content.SharedPreferences;
import android.os.AsyncTask;
import android.os.Bundle;
import android.text.method.PasswordTransformationMethod;
import android.util.Log;
import android.view.Menu;
import android.view.MenuInflater;
import android.view.MenuItem;
import android.view.View;
import android.widget.Button;
import android.widget.CheckBox;
import android.widget.EditText;
import android.widget.TextView;

import com.kaloer.filepicker.FilePickerActivity;

public class AndCryptActivity extends Activity {
	private final static int REQ_SRC_SELECT = 1;
	private final static int REQ_SET_PWD = 2;
	
	public final static String EXTRA_DECRYTION_MODE = "andcrpyt.decrypt_mode";
	
	private final static int DIALOG_STARTUP = 1;
	
	private File srcFile, dstFile;
	
	private EditText srcInput;
	private EditText dstFileName;
	
	private EditText pwdInp;
	
	private CheckBox delInp;
	
	private ProgressDialog progDialog;
	
	private Dialog startupDlg;
	
	private boolean decryptionMode = false;
	
	/** Called when the activity is first created. */
    @Override
    public void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        setContentView(R.layout.main);
        
        if(getIntent().getExtras() != null) {
        	decryptionMode = getIntent().getExtras().getBoolean(EXTRA_DECRYTION_MODE, false);
        }
        
        TextView inputLabel = (TextView)findViewById(R.id.file_input_label);
        
        srcInput = (EditText)findViewById(R.id.file_src);
        dstFileName = (EditText)findViewById(R.id.file_dst_name);
        
        Button openSrc = (Button)findViewById(R.id.file_select_src);
        openSrc.setOnClickListener(new View.OnClickListener() {
			public void onClick(View v) {
				onSelectFile();
			}
		});
        
        Button startBtn = (Button)findViewById(R.id.file_start);
        startBtn.setOnClickListener(new View.OnClickListener() {
			public void onClick(View v) {
				onStartClicked();
			}
		});
        
        delInp = (CheckBox)findViewById(R.id.file_delete_old);
        
        if(decryptionMode) {
        	inputLabel.setText(R.string.select_dec_file);
        	delInp.setText("Delete file after decryption");
        	delInp.setChecked(false);
        }
        else {
        	inputLabel.setText(R.string.select_enc_file);
        	delInp.setText("Delete file after encrpytion");
        	delInp.setChecked(false);
        }
        
        
        if(!getPreferences(0).contains("startup")) {
        	showDialog(DIALOG_STARTUP);
        }
    }
    
    @Override
    protected Dialog onCreateDialog(int id) {
    	if(id == DIALOG_STARTUP) {
    		startupDlg = new Dialog(this);
    		startupDlg.setTitle("Welcome");
    		startupDlg.setContentView(R.layout.startup);
    		
    		Button ok = (Button)startupDlg.findViewById(R.id.startup_ok);
    		ok.setOnClickListener(new View.OnClickListener() {
				public void onClick(View v) {
					onStartDialogOk();
				}
			});
    		
    		return startupDlg;
    	}
    	
    	return null;
    }
    
    private void onStartDialogOk() {
    	CheckBox show = (CheckBox)startupDlg.findViewById(R.id.startup_show);
    	if(show.isChecked()) {
    		SharedPreferences pref = getPreferences(0);
    		SharedPreferences.Editor edit = pref.edit();
    		edit.putBoolean("startup", false);
    		edit.commit();
    	}
    	
    	startupDlg.dismiss();
    	startupDlg = null;
    }
    
    @Override
    public boolean onCreateOptionsMenu(Menu menu) {
        MenuInflater inflater = getMenuInflater();
        inflater.inflate(R.menu.main_menu, menu);
        
        menu.findItem(R.id.menu_dec_mode).setVisible(!decryptionMode);
        menu.findItem(R.id.menu_enc_mode).setVisible(decryptionMode);
        return true;
    }
    
    @Override
    public boolean onOptionsItemSelected(MenuItem item) {
        switch (item.getItemId()) {
            case R.id.menu_dec_mode:
                Intent intent = new Intent(this, AndCryptActivity.class);
                intent.putExtra(EXTRA_DECRYTION_MODE, true);
                finish();
                startActivity(intent);
                return true;
                
            case R.id.menu_enc_mode:
            	finish();
            	startActivity(new Intent(this, AndCryptActivity.class));
                return true;
                
            default:
                return super.onOptionsItemSelected(item);
        }
    }
    
    private void onSelectFile() {
    	Intent intent = new Intent(this, FilePickerActivity.class);
    	startActivityForResult(intent, REQ_SRC_SELECT);
    }
    
    private void onStartClicked() {
    	srcFile = new File(srcInput.getText().toString());
    	
    	if(srcFile == null || !srcFile.exists() || !srcFile.isFile()) {
    		AlertDialog.Builder builder = new AlertDialog.Builder(this);
    		
    		if(decryptionMode) {
    			builder.setMessage(R.string.error_input_dec);
    		}
    		else {
    			builder.setMessage(R.string.error_input_enc);
    		}
    		
    		builder.setNeutralButton("Ok", null);
    		builder.show();
    		return;
    	}
    	
    	String fn = dstFileName.getText().toString();
    	if(fn.length() == 0) {
    		AlertDialog.Builder builder = new AlertDialog.Builder(this);
    		builder.setMessage("Please enter a valid output file name!");
    		builder.setNeutralButton("Ok", null);
    		builder.show();
    		return;
    	}
    	
    	dstFile = new File(srcFile.getParentFile(), fn);
    	
    	//TODO remove after debug
    	//dstFile.delete();
    	
    	if(dstFile.exists()) {
    		AlertDialog.Builder builder = new AlertDialog.Builder(this);
    		builder.setMessage("Output file already exists!");
    		builder.setNeutralButton("Ok", null);
    		builder.show();
    		return;    		
    	}
    	
    	if(decryptionMode) {
    		AlertDialog.Builder d = new AlertDialog.Builder(this);
    		d.setTitle("Enter password");
    		
    		pwdInp = new EditText(this);
    		pwdInp.setTransformationMethod(new PasswordTransformationMethod());
    		
    		d.setView(pwdInp);
    		d.setPositiveButton("Ok", new DialogInterface.OnClickListener() {
				public void onClick(DialogInterface dialog, int which) {
					startDecrypt(srcFile, dstFile, pwdInp.getText().toString());
				}				
			});
    		d.setNegativeButton("Cancel", null);
    		d.show();
    		
    		//startDecrypt(srcFile, dstFile);
    	}
    	else {
    		startActivityForResult(new Intent(this, SetPasswordActivity.class), REQ_SET_PWD);
    		//startEncrypt(srcFile, dst);
    	}
    }
    
    private void startEncrypt(File from, File to, String pwd) {
    	progDialog = new ProgressDialog(this);
    	progDialog.setMessage("Encrypting...");
    	progDialog.setProgressStyle(ProgressDialog.STYLE_HORIZONTAL);
    	progDialog.setProgress(0);
    	progDialog.setMax(100);
    	
    	progDialog.show();
    	
    	EncryptionThread thread = new EncryptionThread();
    	thread.execute(from, to, pwd, delInp.isChecked());
    }

    private void startDecrypt(File from, File to, String pwd) {
    	progDialog = new ProgressDialog(this);
    	progDialog.setMessage("Decrpyting...");
    	progDialog.setProgressStyle(ProgressDialog.STYLE_HORIZONTAL);
    	progDialog.setProgress(0);
    	progDialog.setMax(100);
    	
    	progDialog.show();
    	
    	new DecryptionTask().execute(from, to, pwd, delInp.isChecked());
    }
    
    
    @Override
    protected void onActivityResult(int requestCode, int resultCode, Intent data) {
    	if(resultCode != RESULT_OK) {
    		return;
    	}
    	
    	if(requestCode == REQ_SRC_SELECT) {
    		String path = data.getExtras().getString(FilePickerActivity.EXTRA_FILE_PATH);
    		srcInput.setText(path);
    		
    		srcFile = new File(path);
    		
    		if(decryptionMode) {
    			String fn = srcFile.getName();
    			if(fn.endsWith(".andcrypt")) {
    				fn = fn.substring(0, fn.length() - ".andcrypt".length());
    			}
    			dstFileName.setText(fn);	
    		}
    		else {
    			dstFileName.setText(srcFile.getName() + ".andcrypt");
    		}
    	}
    	else if(requestCode == REQ_SET_PWD) {
    		startEncrypt(srcFile, dstFile, 
    				data.getStringExtra(SetPasswordActivity.EXTRA_PASSWORD));
    	}
    }
    
    private class DecryptionTask extends AsyncTask<Object, Integer, Integer> {

		@Override
		protected Integer doInBackground(Object... params) {
			File from = (File)params[0];
			File to = (File)params[1];
			
			try {
				MessageDigest md = MessageDigest.getInstance(CryptoHelper.HASH_ALG);
				byte[] rawKey = CryptoHelper.passwordToKey(params[2].toString(), CryptoHelper.getStaticSalt());
				SecretKey key = new SecretKeySpec(rawKey, "AES");
				byte[] iv = new byte[CryptoHelper.BLOCK_SIZE];
				Cipher cipher = Cipher.getInstance(CryptoHelper.CRYPT_ALG);
				cipher.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(iv));

				final long totalSize = from.length();
				long done = 0;
				int step = 0;
				
				long written = 0;
				
				
				FileInputStream inp = new FileInputStream(from);
				FileOutputStream out = new FileOutputStream(to);
				
				byte [] buf = new byte[4 * 1024];
				
				while(true) {
					int r = inp.read(buf);
					
					if(r < 0) {
						break;
					}
					done += r;
					
					byte[] dec = cipher.update(buf, 0, r);
					written += dec.length;
					out.write(dec);
					
					step++;
					if(step >= 15) {
						final int percent = (int)(0.75f * (done * 100 / totalSize));
						publishProgress(percent);
						step = 0;
					}
				}
				
				byte[] end = cipher.doFinal();
				written += end.length;
				
				out.write(end);
				
				out.close();
				inp.close();
				
				inp = new FileInputStream(to);
				long readCount = 0;
				int hashSize = md.getDigestLength();
				
				while(readCount < written - hashSize) {
					int r = inp.read(buf, 0, (int)Math.min(buf.length, written - hashSize - readCount));
					
					if(r < 0) {
						break;
					}
					readCount += r;
					
					md.update(buf, 0, r);
				}
				
				byte[] readHash = new byte[hashSize];
				int readHashSize = inp.read(readHash);
				inp.close();
				
				if(readHashSize != hashSize) {
					throw new Exception("File corrupted.");
				}
				
				byte[] calcHash = md.digest();
				
				if(!MessageDigest.isEqual(calcHash, readHash)) {
					throw new Exception("File corrupted.");
				}
				
				RandomAccessFile ra = new RandomAccessFile(to, "rw");
				ra.setLength(written - hashSize);
				ra.close();
				
				if((Boolean)params[3]) {
					from.delete();
				}
				
				return 0;
			}
			catch(Exception e) {
				Log.e("andcrpyt", "Decryption error.", e);
			}
			
			to.delete();
			return 1;
		}
		
		protected void onProgressUpdate(Integer... progress) {
			progDialog.setProgress(progress[0]);
		}	

	    protected void onPostExecute(Integer result) {
	    	progDialog.dismiss();
	    	progDialog = null;
	    	
	        if(result == 0) {
	        	AlertDialog.Builder builder = new AlertDialog.Builder(AndCryptActivity.this);
	        	builder.setMessage("Deryption has been successful.");
	        	builder.setPositiveButton("Ok", null);
	        	builder.show();
	        }
	        else {
	        	AlertDialog.Builder builder = new AlertDialog.Builder(AndCryptActivity.this);
	        	builder.setMessage("Deryption has failed. Wrong password or corrrupted file.");
	        	builder.setPositiveButton("Ok", null);
	        	builder.show();
	        }
	    }
    }
    
    private class EncryptionThread extends AsyncTask<Object, Integer, Integer> {
		@Override
		protected Integer doInBackground(Object... params) {
			File from = (File)params[0];
			File to = (File)params[1];
			String pass = (String)params[2];
			
			try {
				MessageDigest md = MessageDigest.getInstance(CryptoHelper.HASH_ALG);

				byte[] rawKey = CryptoHelper.passwordToKey(pass, CryptoHelper.getStaticSalt());
				SecretKey key = new SecretKeySpec(rawKey, "AES");					
				byte[] iv = new byte[CryptoHelper.BLOCK_SIZE];
				AlgorithmParameterSpec ivParam = new IvParameterSpec(iv);
				Cipher cipher = Cipher.getInstance(CryptoHelper.CRYPT_ALG);
				cipher.init(Cipher.ENCRYPT_MODE, key, ivParam);

				final long totalSize = from.length();
				long done = 0;
				int step = 0;
				
				
				FileInputStream inp = new FileInputStream(from);
				FileOutputStream out = new FileOutputStream(to);
				
				byte [] buf = new byte[4 * 1024];
				
				while(true) {
					int r = inp.read(buf);
					
					if(r < 0) {
						break;
					}
					
					md.update(buf, 0, r);
					byte[] enc = cipher.update(buf, 0, r);

					out.write(enc);
					
					done += r;
					
					step++;
					if(step >= 15) {
						int percent = (int)(done * 100L / totalSize);
						publishProgress(percent);
						step = 0;
					}
				}
				
				byte[] hash = md.digest();
				byte[] enc = cipher.doFinal(hash);
				
				out.write(enc);
				
				inp.close();
				out.close();
			
				if((Boolean)params[3]) {
					from.delete();
				}
				
				return 0;
			}
			catch(Exception e) {
				
			}
			return 1;
		}
		
		protected void onProgressUpdate(Integer... progress) {
			progDialog.setProgress(progress[0]);
		}	

	    protected void onPostExecute(Integer result) {
	    	progDialog.dismiss();
	    	progDialog = null;
	    	
	        if(result == 0) {
	        	AlertDialog.Builder builder = new AlertDialog.Builder(AndCryptActivity.this);
	        	builder.setMessage("Encryption has been successful.");
	        	builder.setPositiveButton("Ok", null);
	        	builder.show();
	        }
	    }
    }
    
}